import requests
import json
import hmac
import base64
from hashlib import sha1
import time
import uuid
import ntplib  # 用于获取网络时间

def get_network_time():
    try:
        ntp_client = ntplib.NTPClient()
        response = ntp_client.request('pool.ntp.org')
        return int(response.tx_time * 1000)
    except Exception as e:
        print(f"获取网络时间失败，使用本地时间: {e}")
        return int(time.time() * 1000)

def make_query(uri, appkey, secret_key):
    timestamp = str(get_network_time())
    signature_nonce = str(uuid.uuid4())
    content = '&'.join((uri, timestamp, signature_nonce))
    digest = hmac.new(secret_key.encode(), content.encode(), sha1).digest()
    sign = base64.urlsafe_b64encode(digest).rstrip(b'=').decode()
    query_params = [
        f'AccessKey={appkey}',
        f'Signature={sign}',
        f'Timestamp={timestamp}',
        f'SignatureNonce={signature_nonce}'
    ]
    return '?' + '&'.join(query_params)

# 提交任务函数
def submit_task(base_url, uri, appkey, appsecret, request_body, retries=3):
    last_error = None
    for attempt in range(retries):
        try:
            full_url = base_url + uri + make_query(uri, appkey, appsecret)
            headers = {"Content-Type": "application/json"}
            print(f"第{attempt+1}次尝试 请求URL: {full_url}")
            response = requests.post(full_url, headers=headers, json=request_body)
            
            # 响应状态码检查
            if response.status_code != 200:
                error_msg = f"HTTP错误 {response.status_code}: {response.text}"
                if response.status_code == 401:
                    error_msg += " (签名验证失败，请检查密钥和时间)"
                last_error = error_msg
                if attempt == retries - 1:
                    raise Exception(f"提交任务失败：{last_error}")
                continue
                
            # JSON解析检查
            try:
                result = response.json()
                print(f"响应解析结果: {result}")
            except Exception as e:
                last_error = f"JSON解析失败：{str(e)} 响应内容：{response.text}"
                if attempt == retries - 1:
                    raise Exception(f"提交任务失败：{last_error}")
                continue

            # generate_uuid有效性检查
            generate_uuid = result.get("data", {}).get("generateUuid", "")
            if not generate_uuid:
                last_error = f"无效UUID 响应结果：{result}"
                if attempt == retries - 1:
                    raise Exception(f"提交任务失败：{last_error}")
                continue
                
            return {
                "success": True,
                "data": result["data"]
            }
            
        except requests.RequestException as e:
            last_error = f"网络请求异常：{str(e)}"
            if attempt == retries - 1:
                raise Exception(f"提交任务失败：{last_error}")
                
    raise Exception(f"提交任务失败：未知错误")

# 查询任务函数，封装轮询逻辑
def query_task(base_url, status_uri, appkey, appsecret, generate_uuid, max_retries=120, retry_interval=5):
    timeout = max_retries * retry_interval
    start_time = time.time()
    retries = 0

    while retries < max_retries:
        current_time = time.time()
        if (current_time - start_time) > timeout:
            raise Exception(f"查询任务失败：{timeout}s任务超时，已退出轮询")

        full_url = base_url + status_uri + make_query(status_uri, appkey, appsecret)
        headers = {"Content-Type": "application/json"}
        status_data = {"generateUuid": generate_uuid}
        
        try:
            status_response = requests.post(full_url, headers=headers, data=json.dumps(status_data))
            if status_response.status_code != 200:
                raise Exception(f"状态码异常[{status_response.status_code}] 响应内容：{status_response.text}")
                
            status_result = status_response.json()
            
            if status_result and status_result.get('data', {}).get('images'):
                for url in status_result['data']['images']:
                    if url and url.get("imageUrl"):
                        return {
                            "success": True,
                            "imageUrl": url["imageUrl"],
                            "message": "任务完成，获取到图像数据"
                        }
            
            print(f"任务尚未完成，等待 {retry_interval} 秒...")
            time.sleep(retry_interval)
            retries += 1
            
        except ValueError as e:
            raise Exception(f"查询任务失败：JSON解析异常 - {str(e)}")
        except requests.RequestException as e:
            if retries == max_retries - 1:
                raise Exception(f"查询任务失败：网络请求异常 - {str(e)}")
            continue
            
    raise Exception(f"查询任务失败：达到最大重试次数 {max_retries}")

def main(prompt: str, appkey: str, appsecret: str, model_type: str = "liblibai") -> dict | str:
    try:
        if not appkey or not appsecret:
            raise ValueError("appkey 和 appsecret 不能为空，请检查输入。")

        local_time = int(time.time() * 1000)
        network_time = get_network_time()
        time_diff = abs(local_time - network_time)
        if time_diff > 3000:
            print(f"警告：本地时间和网络时间偏差 {time_diff} 毫秒，请检查系统时间。")

        base_url = "https://openapi.liblibai.cloud"

        ##################################3.1.1 星流Star-3 Alpha文生图#################################################
        generate_uri = '/api/generate/webui/text2img/ultra'
        status_uri = '/api/generate/webui/status'

        if model_type == "liblibai":
            request_body = {
                "templateUuid": "5d7e67009b344550bc1aa6ccbfa1d7f4", #3af36dd5a61e4da88c6cb5eb57a8fe2e
                "generateParams": {
                    "prompt": prompt,
                    "aspectRatio": "portrait",
                    "imageSize": {
                        "width": 768,
                        "height": 1024
                    },
                    "imgCount": 1,
                    "steps": 30
                }
            }
        else:
            raise ValueError(f"不支持的模型类型: {model_type}")
        ##################################3.1.1 星流Star-3 Alpha文生图#################################################

        ##################################3.1.2 星流Star-3 Alpha图生图#################################################
        # generate_uri = "/api/generate/webui/img2img/ultra"
        # status_uri = "/api/generate/webui/status"
        # request_body = {
        #     "templateUuid": "07e00af4fc464c7ab55ff906f8acf1b7",
        #     "generateParams": {
        #         "checkPointId": "0ea388c7eb854be3ba3c6f65aac6bfd3",
        #         "prompt": "1 girl wear sunglasses",
        #         "clipSkip": 2,
        #         "sampler": 15,
        #         "steps": 20,
        #         "cfgScale": 7,
        #         "randnSource": 0,
        #         "seed": -1,
        #         "imgCount": 1,
        #         "restoreFaces": 0,
        #         "sourceImage": "https://liblibai-online.liblib.cloud/img/081e9f07d9bd4c2ba090efde163518f9/7c1cc38e-522c-43fe-aca9-07d5420d743e.png",
        #         "resizeMode": 0,
        #         "resizedWidth": 1024,
        #         "resizedHeight": 1536,
        #         "mode": 4,
        #         "denoisingStrength": 0.75,
        #         "additionalNetwork": [
        #             {
        #                 "modelId": "31360f2f031b4ff6b589412a52713fcf",
        #                 "weight": 0.3
        #             },
        #             {
        #                 "modelId": "365e700254dd40bbb90d5e78c152ec7f",
        #                 "weight": 0.6
        #             }
        #         ],
        #         "inpaintParam": {
        #             "maskImage": "https://liblibai-online.liblib.cloud/img/081e9f07d9bd4c2ba090efde163518f9/323fc358-618b-4c7d-b431-7d890209e5a5.png",
        #             "maskBlur": 4,
        #             "maskPadding": 32,
        #             "maskMode": 0,
        #             "inpaintArea": 0,
        #             "inpaintingFill": 1
        #         },
        #         "controlNet": [
        #             {
        #                 "unitOrder": 1,
        #                 "sourceImage": "https://liblibai-online.liblib.cloud/img/081e9f07d9bd4c2ba090efde163518f9/7c1cc38e-522c-43fe-aca9-07d5420d743e.png",
        #                 "width": 1024,
        #                 "height": 1536,
        #                 "preprocessor": 3,
        #                 "annotationParameters": {
        #                     "depthLeres": {
        #                         "preprocessorResolution": 1024,
        #                         "removeNear": 0,
        #                         "removeBackground": 0
        #                     }
        #                 },
        #                 "model": "6349e9dae8814084bd9c1585d335c24c",
        #                 "controlWeight": 1,
        #                 "startingControlStep": 0,
        #                 "endingControlStep": 1,
        #                 "pixelPerfect": 1,
        #                 "controlMode": 0,
        #                 "resizeMode": 1,
        #                 "maskImage": ""
        #             }
        #         ]
        #     }
        # }
        ##################################3.1.2 星流Star-3 Alpha图生图#################################################

        ##################################6.1 ComfyUI工作流生图#################################################
        # generate_uri = "/api/generate/comfyui/app"
        # status_uri = "/api/generate/comfy/status"
        # request_body ={
        #     "templateUuid": "4df2efa0f18d46dc9758803e478eb51c",
        #     "generateParams": {
        #         "12": {
        #             "class_type": "LoadImage",
        #             "inputs": {
        #                 "image": "https://liblibai-tmp-image.liblib.cloud/img/baf2e419ce1cb06812314957efd2e067/af0c523d3d2b4092ab45c64c72e4deb76babb12e9b8a178eb524143c3b71bf85.png"
        #             }
        #         },
        #         "112": {
        #             "class_type": "ImageScale",
        #             "inputs": {
        #                 "width": 768
        #             }
        #         },
        #         "136": {
        #             "class_type": "RepeatLatentBatch",
        #             "inputs": {
        #                 "amount": 4
        #             }
        #         },
        #         "137": {
        #             "class_type": "LatentUpscaleBy",
        #             "inputs": {
        #                 "scale_by": 1.5
        #             }
        #         },
        #         "workflowUuid": "2f22ab7ce4c044afb6d5eee2e61547f3"
        #     }
        # }
        ##################################6.1 ComfyUI工作流生图#################################################

        # 提交生成图像任务
        generate_response = submit_task(base_url, generate_uri, appkey, appsecret, request_body)

        # 直接获取generate_uuid（因submit_task已校验）
        generate_uuid = generate_response["data"]["generateUuid"]

        # 调用查询任务函数
        return query_task(base_url, status_uri, appkey, appsecret, generate_uuid)
    
    except Exception as e:
        return {
            "success": False,
            "imageUrl": "",
            "message": str(e)
        }